from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, roc_auc_score, classification_report, balanced_accuracy_score
from imblearn.over_sampling import SMOTE
from xgboost import XGBClassifier
from lightgbm import LGBMClassifier
import warnings
#  删除了本地定义的 data_preprocessing 函数
#  改为从 utils.common 导入
from utils.common import data_preprocessing  # ✅ 引入统一预处理函数
from utils.feature_engineering import feature_engineering
# 绘图
from sklearn.metrics import roc_curve, auc
import matplotlib.pyplot as plt
#导出模型
import joblib
import os
warnings.filterwarnings('ignore')

plt.rcParams['font.family'] = 'SimHei'
plt.rcParams['font.size'] = 15
plt.rcParams['axes.unicode_minus'] = False

import pandas as pd
import matplotlib.pyplot as plt

def plot_roc_curve(y_true, y_score, model_name, save_path=None):
    """
    绘制单个模型的 ROC 曲线
    :param y_true: 真实标签
    :param y_score: 模型预测概率
    :param model_name: 模型名称
    :param save_path: 保存路径，如果为 None 则不保存
    """
    fpr, tpr, _ = roc_curve(y_true, y_score)
    roc_auc = auc(fpr, tpr)

    plt.figure(figsize=(8, 6))
    plt.plot(fpr, tpr, color='darkorange', lw=2, label=f'ROC curve (AUC = {roc_auc:.4f})')
    plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('False Positive Rate', fontsize=12)
    plt.ylabel('True Positive Rate', fontsize=12)
    plt.title(f'Receiver Operating Characteristic - {model_name}', fontsize=14)
    plt.legend(loc="lower right", fontsize=12)
    plt.grid(True)
    plt.tight_layout()

    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
    plt.show()

def train_model(train_data, test_data):
    # 特征和标签
    X_train = train_data.iloc[:, 1:]
    y_train = train_data.iloc[:, 0]
    X_test = test_data.iloc[:, 1:]
    y_test = test_data.iloc[:, 0]

    # 过采样
    smote = SMOTE(sampling_strategy=0.4, random_state=42)
    X_resample, y_resample = smote.fit_resample(X_train, y_train)

    # 划分数据集
    X_train_, X_val_, y_train_, y_val_ = train_test_split(
        X_resample, y_resample,
        test_size=0.2,
        random_state=42,
        stratify=y_resample
    )

    # 标准化
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train_)
    X_val_scaled = scaler.transform(X_val_)
    X_test_scaled = scaler.transform(X_test)

    # 模型定义
    models = {
        "XGBClassifier": XGBClassifier(
            objective='binary:logistic',
            eval_metric='auc',
            learning_rate=0.1,
            max_depth=3,
            subsample=0.8,
            colsample_bytree=0.7,
            n_estimators=201,
            random_state=1,
            reg_lambda=26,
            use_label_encoder=False),
        "LGBMClassifier": LGBMClassifier(
            max_depth=11,
            num_leaves=30,
            min_child_samples=18,
            bagging_fraction=0.6,
            feature_fraction=0.4,
            random_state=26,
            verbose=-1)
    }

    # 模型训练与评估
    for name, model in models.items():
        print(f"【{name}】")
        model.fit(X_train_scaled, y_train_)
        y_test_pred = model.predict(X_test_scaled)
        y_pred_prob = model.predict_proba(X_test_scaled)[:, 1]

        # 绘图
        plot_roc_curve(y_test, y_pred_prob, model_name=name, save_path=f'../results/roc_curve_{name}.png')

        # 打印评估结果
        print(f"Test Accuracy: {accuracy_score(y_test, y_test_pred):.4f}")
        print(f"Balanced Accuracy: {balanced_accuracy_score(y_test, y_test_pred):.4f}")
        print(f"ROC AUC Score: {roc_auc_score(y_test, y_pred_prob):.4f}")
        print("Classification Report:")
        print(classification_report(y_test, y_test_pred))
        # 保存模型
        model_save_path = os.path.join('../model', f'{name}.joblib')
        joblib.dump(model, model_save_path)
        print(f"Successful:模型 [{name}] 已保存至：{model_save_path}")
        print("-" * 60)
        #
        # # 特征重要性
        # if name == "LGBMClassifier":
        #     feature_importances = model.feature_importances_
        #     feature_names = X_train.columns
        #     importance_df = pd.DataFrame({
        #         'Feature': feature_names,
        #         'Importance': feature_importances
        #     }).sort_values(by='Importance', ascending=False)
        #
        #     # 输出特征重要性
        #     print("LGBMClassifier 特征重要性：")
        #     print(importance_df)
        #
        #     # 绘制优化后的条形图
        #     plt.figure(figsize=(16, 12))
        #     bars = plt.barh(importance_df['Feature'], importance_df['Importance'], color='skyblue')
        #
        #     # 添加数值标签
        #     for bar in bars:
        #         width = bar.get_width()
        #         plt.text(width + 0.001, bar.get_y() + bar.get_height() / 2, f'{width:.4f}', va='center', fontsize=10)
        #
        #     plt.xlabel('Feature Importance', fontsize=14)
        #     plt.ylabel('Feature', fontsize=14)
        #     plt.title('lightgbm Feature Importance (Sorted)', fontsize=16)
        #     plt.grid(axis='x', linestyle='--', alpha=0.7)
        #     plt.tight_layout()
        #
        #     # 调整字体和中文支持
        #     plt.rcParams['font.family'] = 'SimHei'
        #     plt.gca().invert_yaxis()  # 重要性高的特征在上
        #
        #     # 保存并显示
        #     plt.savefig('../results/lightgbm_feature_importance_optimized.png', dpi=300, bbox_inches='tight')
        #     plt.show()

# if __name__ == '__main__':
#     # 加载训练集并生成编码器
#     train_data, encoders = data_preprocessing(r'../data/train.csv', is_train=True)
#
#     # 使用训练集的编码器预处理测试集
#     test_data, _ = data_preprocessing(r'../data/test.csv', encoders=encoders, is_train=False)
#
#     # 保存处理后的测试集用于后续测试
#     test_data.to_csv('../data/test_processed.csv', index=False)
#
#     # 特征工程
#     engineered_train_data = feature_engineering(train_data)
#     engineered_test_data = feature_engineering(test_data)
#
#     # 模型训练与评估
#     train_model(engineered_train_data, engineered_test_data)

def main():
    # 加载训练集并生成编码器
    train_data, encoders = data_preprocessing(r'../data/train.csv', is_train=True)

    # 使用训练集的编码器预处理测试集
    test_data, _ = data_preprocessing(r'../data/test.csv', encoders=encoders, is_train=False)

    # 保存处理后的测试集用于后续测试
    test_data.to_csv('../data/test_processed.csv', index=False)

    # 特征工程
    engineered_train_data = feature_engineering(train_data)
    engineered_test_data = feature_engineering(test_data)

    # 模型训练与评估
    train_model(engineered_train_data, engineered_test_data)

if __name__ == '__main__':
    main()